import numpy as np
import faiss

# 模拟数据
np.random.seed(123)
d = 128  # 向量维度
nb = 100000  # 数据库向量数量
nq = 10  # 查询向量数量

# 随机生成数据库和查询向量
xb = np.random.random((nb, d)).astype('float32')
xq = np.random.random((nq, d)).astype('float32')

# 1. 进行聚类
n_clusters = int(nb * 0.01)  # 聚类中心的数量（约为向量数量的1%）
kmeans = faiss.Kmeans(d, n_clusters, niter=20, verbose=False)
kmeans.train(xb)
cluster_centroids = kmeans.centroids

# 2. 建立聚类中心的图（这里使用faiss的IndexFlatL2进行暴力图）
index_flat = faiss.IndexFlatL2(d)
index_flat.add(cluster_centroids)

# 3. 检索时首先在聚类中心进行查询
D, I = index_flat.search(xq, 1)  # 在聚类中心中找到最近的一个

# 4. 进一步在对应的聚类中进行精确匹配
final_results = []
for i in range(nq):
    cluster_id = I[i][0]
    
    # 找出属于该聚类的所有点
    cluster_mask = kmeans.index.search(xb, 1)[1].flatten() == cluster_id
    cluster_vectors = xb[cluster_mask]

    # 在该聚类的点中进行匹配
    cluster_index = faiss.IndexFlatL2(d)
    cluster_index.add(cluster_vectors)
    
    D_cluster, I_cluster = cluster_index.search(xq[i:i+1], 5)  # 取前5个最近邻
    final_results.append((D_cluster, I_cluster))

# 打印结果
for i, result in enumerate(final_results):
    print(f"Query {i}: Nearest neighbors in cluster: {result[1]}")
